import scipy
import numpy as np
import matplotlib.pyplot as plt
import librosa

# y是音频数据，sr是采样率

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号



def draw_FFT(Path):  # 传入音乐路径，进行快读傅里叶变换，并绘制图像
    print("----------开始快速傅里叶变换----------")
    signal, sample_rate = librosa.load(Path,sr=None)  # 返回音频文件的采样率和numpy数组（包括音频文件的音频数据）
    # signal其实就是对这段音频文件采样得到的数字点
    print(f"采样率为{sample_rate}.{signal.shape}")

    # -----------------进行FFT变换-----------------
    spectrum = scipy.fft.fft(signal[:44000])  # 通过快速傅里叶变换，计算信号的频谱。添加频率轴
    # 返回复数的矩阵

    freq_axis = np.fft.fftfreq(len(spectrum), d=1 / sample_rate)
    freq_axis = freq_axis[:len(freq_axis) // 2]  # 因为对称，所以去一半

    Amplitude = np.abs(spectrum[:len(freq_axis)])  # abs计算复数的模
    Amplitude_nor = Amplitude / len(freq_axis)  # 归一化

    # -----------------绘制频谱图----------------

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
    ax1.plot(freq_axis, Amplitude)
    ax1.set_xlabel("频率")
    ax1.set_ylabel("幅度")

    ax2.plot(freq_axis, Amplitude_nor)
    ax2.set_xlabel("频率")
    ax2.set_ylabel("归一化幅度")
    plt.show()

    print("----------完成快速傅里叶变换----------")
    return freq_axis, Amplitude  # 返回频率轴和对应幅度


if __name__ == "__main__":
    path = "../music sample/花海伴奏节选.mp3"
    a, b = draw_FFT(path)
